from zhipuai import ZhipuAI


class Agent_role:

    def __init__(self):
        self.message_norm = []
        self.message_architecture = []
        self.message_answer = []
        self.tools = []
        self.message_expanding = []

    def zhipu_api(self, m, k = "", model = 'glm-4', tools = None):  # 填写自己的key
        client = ZhipuAI(api_key = k)
        response = client.chat.completions.create(
            model = model,
            messages = m,
            # stream = True
            tools = tools,
            max_tokens = 8192,  # max:8192
            temperature = 0.97
        )
        # #异步调用
        # task_id = response.id
        # task_status = ''
        # get_cnt = 0
        # while task_status != 'SUCCESS' and task_status != 'FAILED' and get_cnt <= 40:
        #     result_response = client.chat.asyncCompletions.retrieve_completion_result(id = task_id)
        #     if 'choices' in result_response:
        #         return result_response.choices[0].message
        #     task_status = result_response.task_status
        #     time.sleep(2)
        #     get_cnt += 1
        return response.choices[0].message

    def mesage_template(self, role: str = 'user', content: str = '') -> dict:
        return {"role": role, "content": content}

    def tool_template(self, name: str, description: str, parameters: dict):
        tool = {
            "type": "function",
            "function": {
                "name": name,
                "description": description,
                "parameters": parameters
            }
        }

    def parameters_template(self, properties):
        parameters = {
            "type": "object",
            "properties": properties,
            "required": list(properties.keys())
        }

    def Tools(self):
        properties = {"keyword": {"description": "关键词", "type": "string"}}
        name = "knowledge_learning"
        description = "若你不确定，请根据关键词,调用搜索引擎学习相应的背景知识"
        self.tools.append(self.tool_template(name, description, self.parameters_template(properties)))

    def problem_norm(self, question):
        c = '你作为一名经验丰富的语言专家，你擅长理解和分析语句，对问题语句进行拆分。' \
            '例如：“请帮我编写一篇关于爱情的长篇小说。”拆分成“1.编写关于爱情的小说；“' \
            '不要输出其他无关语句，直接得到改写结果。下面是我的问题：' + question
        self.message_norm.append(self.mesage_template(content = c))
        m = self.zhipu_api(self.message_norm)
        self.message_norm.append(m)
        return m.content

    def problem_architecture(self, norm_question):
        c = '你作为一名作家，你了解不同故事和小说的基本剧情，擅长撰写不同类型的小说基本剧情。不考虑字数的要求，只生成小说基本剧情，以“第一章”、”第二章“的形式划分章节。' \
            '各种小说角色需要具备唯一名字。小说应具备人物形象、故事情节、典型环境等要素。不要输出其他无关语句，直接输出小说的基本剧情，不要提出建议，最后不要总结小说。'
        self.message_architecture.append(self.mesage_template(content = c))
        self.message_architecture.append(self.mesage_template(content = norm_question))
        m = self.zhipu_api(self.message_architecture, tools = self.Tools())
        self.message_architecture.append(m)
        return m.content

    def problem_answer(self, norm_question, architecture_question, example = None):
        c = '你是一名优秀的网络小说作家，擅长撰写不同类型的小说。根据小说编写要求和小说基本剧情，对小说基本剧情进行扩写并达到字数要求。' \
            '不要产生新角色，不要输出其他无关语句，直接输出所撰写的小说，不要自己总结。'
        self.message_answer.append(self.mesage_template(content = c))
        self.message_answer.append(self.mesage_template(content = "小说编写要求：" + norm_question + '\n'
                                                                                              "小说基本剧情：" + architecture_question))
        m = self.zhipu_api(self.message_answer)
        self.message_answer.append(m)
        return m.content

    def expanding(self, above, below, word_len):
        c = "你是一名优秀的网络小说作家，擅长结合上文，在不产生新的实体的情况下对给出的下文进行字数扩充。" \
            "例如：上文：“李雷发现了一个入口”，下文：“李雷穿越到了异世界”，" \
            "将下文扩充为“李雷小心翼翼地走进这个未知的入口，突然眼前一黑，在通过一阵头昏目眩后，渐渐失去了意识。等到李雷重新恢复意识时，" \
            "他发现自己躺在一片柔软的草地上，身边繁花似锦，天空中飘着五彩斑斓的云朵。他的眼睛被这片美景深深吸引，感觉自己仿佛进入了一个童话世界。" \
            "他站起身来，打量四周，发现自己身处在一个神秘的森林里。李雷环顾四周，试图寻找一些熟悉的事物。但他很快就失望了，这里的一切都和他所知道的世界大相径庭。”。" \
            "以下是我给出的上下文，上文：“" + str(above) + "”；下文：“" + str(below) + "”。对下文进行不少于" + str(
            word_len) + "字的扩充，不要自己补充结局，只在给出的下文上进行字数扩充。" \
                        "不要产生新角色，不要输出其他无关语句，直接输出扩充字数后的下文，不要自己总结。"
        m = self.zhipu_api([self.mesage_template(content = c)])
        self.message_expanding.append(m)
        return m.content
